package com.cloudea.learn.optimizer;

import com.cloudea.learn.Tensor;

/**
 * Adam优化器
 * @author cloudea
 * @date 2022/01/03
 */
public class Adam implements Optimizer {
    private final double alpha;
    private final double beta1;
    private final double beta2;
    private final double epsilon;
    private Tensor m;
    private Tensor v;
    private int t = 0;

    /**
     * 构造函数
     * @param alpha 学习率，通常是0.001
     * @param beta1 通常是0.9
     * @param beta2 通常是0.999
     * @param epsilon 通常是1e-8
     */
    public Adam(double alpha, double beta1, double beta2, double epsilon){
        this.alpha = alpha;
        this.beta1 = beta1;
        this.beta2 = beta2;
        this.epsilon = epsilon;
    }

    @Override
    public void optimize(Tensor param, Tensor grid) {
        if(m == null){ m = new Tensor(param.getShape());}
        if(v == null){ v = new Tensor(param.getShape()); }
        t = t + 1;
        for (int i = 0; i < param.getSize(); i++){
            m.setByLocation(i, beta1 * m.getByLocation(i) + (1 - beta1) * grid.getByLocation(i));
            v.setByLocation(i, beta2 * v.getByLocation(i) + (1 - beta2) * Math.pow(grid.getByLocation(i), 2));
        }
        for (int i = 0; i < param.getSize(); i++){
            double mHat = m.getByLocation(i) / (1 - Math.pow(beta1, t));
            double vHat = v.getByLocation(i) / (1 - Math.pow(beta2, t));
            //更新参数
            param.setByLocation(i, param.getByLocation(i) - alpha * mHat/ (Math.sqrt(vHat) + epsilon));
        }
    }

    /**
     * Adam的建造者
     */
    public static class Builder implements OptimizerBuilder{

        private double alpha = 0.001;
        private double beta1 = 0.9;
        private double beta2 = 0.999;
        private double epsilon = 1e-8;


        public Builder setAlpha(double alpha) {
            this.alpha = alpha;
            return this;
        }


        public Builder setBeta1(double beta1) {
            this.beta1 = beta1;
            return this;
        }


        public Builder setBeta2(double beta2) {
            this.beta2 = beta2;
            return this;
        }


        public Builder setEpsilon(double epsilon) {
            this.epsilon = epsilon;
            return this;
        }

        @Override
        public Optimizer build() {
            return new Adam(alpha, beta1, beta2, epsilon);
        }
    }
}
